#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
姿态检测模块
作者: 元歌
开发地: 广州南沙铁建中心
时间: 2024年10月
"""

import cv2
import mediapipe as mp
import numpy as np
from typing import List, Tuple, Optional

class PoseDetector:
    """姿态检测器，用于检测人体骨骼关键点"""
    
    def __init__(self):
        self.mp_pose = mp.solutions.pose
        self.pose = self.mp_pose.Pose(
            static_image_mode=False,
            model_complexity=1,
            enable_segmentation=False,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5
        )
        self.mp_drawing = mp.solutions.drawing_utils
        
    def detect_pose(self, image: np.ndarray) -> Tuple[Optional[object], np.ndarray]:
        """
        检测图像中的人体姿态
        
        Args:
            image: 输入图像 (BGR格式)
            
        Returns:
            pose_landmarks: 姿态关键点，如果没有检测到则返回None
            annotated_image: 标注了关键点的图像
        """
        # 转换颜色空间
        rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 检测姿态
        results = self.pose.process(rgb_image)
        
        # 创建标注图像
        annotated_image = image.copy()
        if results.pose_landmarks:
            # 绘制骨架与关键点编号
            h, w = annotated_image.shape[:2]
            point_color = (0, 255, 255)
            line_color = (0, 200, 255)
            text_color = (255, 255, 255)
            font = cv2.FONT_HERSHEY_SIMPLEX
            # 字号略调小
            scale_base = max(0.3, min(0.6, min(h, w) / 400.0))
            radius = max(2, int(3 * scale_base))
            thickness = max(1, int(2 * scale_base))
            landmarks = results.pose_landmarks.landmark
            try:
                for conn in self.mp_pose.POSE_CONNECTIONS:
                    i, j = conn
                    li = landmarks[i]
                    lj = landmarks[j]
                    xi, yi = int(li.x * w), int(li.y * h)
                    xj, yj = int(lj.x * w), int(lj.y * h)
                    if 0 <= xi < w and 0 <= yi < h and 0 <= xj < w and 0 <= yj < h:
                        cv2.line(annotated_image, (xi, yi), (xj, yj), line_color, thickness)
            except Exception:
                pass
            for idx, lm in enumerate(landmarks):
                x_px, y_px = int(lm.x * w), int(lm.y * h)
                if x_px < 0 or y_px < 0 or x_px >= w or y_px >= h:
                    continue
                cv2.circle(annotated_image, (x_px, y_px), radius, point_color, -1)
                cv2.putText(annotated_image, str(idx), (x_px + 3, y_px - 3), font, scale_base, text_color, 1, cv2.LINE_AA)
        
        return results.pose_landmarks, annotated_image
    
    def get_landmark_coordinates(self, landmarks) -> List[Tuple[float, float, float]]:
        """
        获取关键点坐标
        
        Args:
            landmarks: 姿态关键点对象
            
        Returns:
            关键点坐标列表 [(x, y, z), ...]
        """
        if not landmarks:
            return []
        
        coords = []
        for landmark in landmarks.landmark:
            coords.append((landmark.x, landmark.y, landmark.z))
        
        return coords
    
    def get_landmark_by_name(self, landmarks, landmark_name: str) -> Optional[Tuple[float, float, float]]:
        """
        根据名称获取特定关键点坐标
        
        Args:
            landmarks: 姿态关键点对象
            landmark_name: 关键点名称
            
        Returns:
            关键点坐标 (x, y, z) 或 None
        """
        if not landmarks:
            return None
        
        landmark_names = {
            'nose': 0,
            'left_eye_inner': 1, 'left_eye': 2, 'left_eye_outer': 3,
            'right_eye_inner': 4, 'right_eye': 5, 'right_eye_outer': 6,
            'left_ear': 7, 'right_ear': 8,
            'mouth_left': 9, 'mouth_right': 10,
            'left_shoulder': 11, 'right_shoulder': 12,
            'left_elbow': 13, 'right_elbow': 14,
            'left_wrist': 15, 'right_wrist': 16,
            'left_pinky': 17, 'right_pinky': 18,
            'left_index': 19, 'right_index': 20,
            'left_thumb': 21, 'right_thumb': 22,
            'left_hip': 23, 'right_hip': 24,
            'left_knee': 25, 'right_knee': 26,
            'left_ankle': 27, 'right_ankle': 28,
            'left_heel': 29, 'right_heel': 30,
            'left_foot_index': 31, 'right_foot_index': 32
        }
        
        if landmark_name not in landmark_names:
            return None
        
        landmark_idx = landmark_names[landmark_name]
        landmark = landmarks.landmark[landmark_idx]
        return (landmark.x, landmark.y, landmark.z)
    
    def calculate_angle(self, point1: Tuple[float, float, float], 
                       point2: Tuple[float, float, float], 
                       point3: Tuple[float, float, float]) -> float:
        """
        计算三点之间的角度
        
        Args:
            point1, point2, point3: 三个点的坐标 (x, y, z)
            
        Returns:
            角度（度）
        """
        # 计算向量
        vector1 = np.array([point1[0] - point2[0], point1[1] - point2[1]])
        vector2 = np.array([point3[0] - point2[0], point3[1] - point2[1]])
        
        # 计算角度
        cos_angle = np.dot(vector1, vector2) / (np.linalg.norm(vector1) * np.linalg.norm(vector2))
        cos_angle = np.clip(cos_angle, -1.0, 1.0)  # 防止数值误差
        angle = np.arccos(cos_angle)
        
        return np.degrees(angle)
    
    def calculate_distance(self, point1: Tuple[float, float, float], 
                          point2: Tuple[float, float, float]) -> float:
        """
        计算两点之间的距离
        
        Args:
            point1, point2: 两个点的坐标 (x, y, z)
            
        Returns:
            距离
        """
        return np.sqrt((point1[0] - point2[0])**2 + (point1[1] - point2[1])**2 + (point1[2] - point2[2])**2)


